#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#pragma OPENCL EXTENSION cl_arm_printf : enable

__kernel void conv_1x1_naive(__global float4 *output_buffer,                // 32个36个14400行
                             __global float4 *input_buffer,
                             __global float16 *weight_buffer,               // 32*36*32*2*4*4
                             int4 input_sliced_shape, int4 weight_sliced_shape) {    // 1, 36, 14400, 32;   64, 6, 6, 32
    int dst_x = get_global_id(0) * 2;
    int dst_y = get_global_id(1);
    int dst_z = get_global_id(2) * 2;

    int src_width = input_sliced_shape.z, src_height = input_sliced_shape.y, src_slice = input_sliced_shape.w;
    int weight_o_slice = weight_sliced_shape.x, weight_i_slice = weight_sliced_shape.w;
    int weight_height = weight_sliced_shape.y, weight_width = weight_sliced_shape.z; 
    int weight_kernel_size = weight_height * weight_width;
    int dst_width = src_width, dst_height = src_height, dst_slice = weight_o_slice;

    if (dst_x >= dst_width || dst_y >= dst_height || dst_z >= dst_slice) return;

    // int x_idx = min(dst_x + 0, src_width-1);
    // int y_idx = min(dst_y, src_height - 1);
    int input_idx = dst_y * src_width + dst_x;     // start_idx

    __global float16 *weight_start = weight_buffer + ((dst_z/2)*weight_kernel_size+dst_y)*weight_i_slice*2;
    float4 group0_v0 = 0, group0_v1 = 0, group1_v0 = 0, group1_v1 = 0;
    for (int ic=0; ic<src_slice; ++ic) {
        float4 input0 = input_buffer[input_idx], input1 = input_buffer[input_idx + 1];
        float16 weight0 = weight_start[0], weight1 = weight_start[1];
        group0_v0 += weight0.s0123 * input0.s0;                           // weight0.s0123是4个filter中的channel_s同一个位置, s0.s0是input中channel_s上同一个位置的值，这样子算出oc4
        group0_v0 += weight0.s4567 * input0.s1;                           // weight0.s4567是4个filter中的channel_s+1的同一个位置，这里是沿着ic维度的累加
        group0_v0 += weight0.s89ab * input0.s2;
        group0_v0 += weight0.scdef * input0.s3;                           // 直到channel_s+3，到此一个ic4计算完成
        group0_v1 += weight0.s0123 * input1.s0;                           // 算input中的下一个ic4
        group0_v1 += weight0.s4567 * input1.s1;
        group0_v1 += weight0.s89ab * input1.s2;
        group0_v1 += weight0.scdef * input1.s3;

        group1_v0 += weight1.s0123 * input0.s0;
        group1_v0 += weight1.s4567 * input0.s1;
        group1_v0 += weight1.s89ab * input0.s2;
        group1_v0 += weight1.scdef * input0.s3;
        group1_v1 += weight1.s0123 * input1.s0;
        group1_v1 += weight1.s4567 * input1.s1;
        group1_v1 += weight1.s89ab * input1.s2;
        group1_v1 += weight1.scdef * input1.s3;

        input_idx = input_idx + src_width * src_height;
        weight_start = weight_start + 2;
    }

    // if (dst_x + 0 < dst_width && dst_y + 0 < dst_height) {
        output_buffer[(dst_z * dst_height + dst_y) * dst_width + dst_x] = group0_v0;
    // }
    // if (dst_x + 1 < dst_width && dst_y + 0 < dst_height) {
        output_buffer[(dst_z * dst_height + dst_y) * dst_width + dst_x + 1] = group0_v1;
    // }
    // if (dst_z + 1 >= dst_slice) 
        // return;
    // if (dst_x + 0 < dst_width && dst_y + 0 < dst_height) {
        output_buffer[((dst_z + 1) * dst_height + dst_y) * dst_width + dst_x] = group1_v0;
    // }
    // if (dst_x + 1 < dst_width && dst_y + 0 < dst_height) {
        output_buffer[((dst_z + 1) * dst_height + dst_y + 0) * dst_width + dst_x + 1] = group1_v1;
    // }
}